package tech.realcpf.duckengine1.integration.adbc;

import org.apache.arrow.adbc.core.AdbcConnection;
import org.apache.arrow.adbc.core.AdbcDatabase;
import org.apache.arrow.adbc.core.AdbcStatement;
import org.apache.arrow.adbc.driver.jdbc.JdbcDriver;
import org.apache.arrow.c.ArrowArrayStream;
import org.apache.arrow.c.Data;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.ipc.ArrowReader;
import org.duckdb.DuckDBConnection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.DriverManager;
import java.sql.Statement;
import java.util.Map;
import java.util.UUID;

public class CommonAdbcReader {
    private static final Logger LOGGER = LoggerFactory.getLogger(CommonAdbcReader.class);

    private final Map<String, Object> parameters;
    private String tmpTableName;
    public CommonAdbcReader(Map<String,Object> parameters) {
        this.parameters = parameters;
    }
    public String tmpTableName() {
        tmpTableName =  "dual_" + UUID.randomUUID().toString().replace("-","");
        return tmpTableName;
    }
    public void loadBySql(String sourceSql, String targetSql, String init) {
        try(BufferAllocator allocator = new RootAllocator();
            AdbcDatabase db = new JdbcDriver(allocator).open(parameters);
            AdbcConnection adbcConnection = db.connect();
            AdbcStatement stmt = adbcConnection.createStatement()) {
            stmt.setSqlQuery(sourceSql);
            AdbcStatement.QueryResult queryResult = stmt.executeQuery();
            ArrowReader reader = queryResult.getReader();
            try(ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) {
                Data.exportArrayStream(allocator,reader,stream);
                try(DuckDBConnection conn =  (DuckDBConnection) DriverManager.getConnection("jdbc:duckdb:")) {
                    conn.registerArrowStream(tmpTableName,stream);
                    try(Statement statement = conn.createStatement()) {
                        statement.execute(init);
                        statement.execute(targetSql);
                    }
                }

            }catch (Exception e) {
                LOGGER.error("write data error",e);
            }
        }catch (Exception e) {
            LOGGER.error("load data error",e);
        }
    }
}
